# 导入操作系统相关功能模块
import os
# 导入数值计算库
import numpy as np
# 从 huggingface 库导入 AutoModel 和 AutoTokenizer
from transformers import AutoModel, AutoTokenizer
# 新增导入千问模型相关库
from transformers import AutoModelForCausalLM
# 从 sklearn 库导入 KMeans 聚类算法
from sklearn.cluster import KMeans
# 导入 openpyxl 库用于操作 Excel 文件
import openpyxl
from tqdm import tqdm
import torch

# 定义全局变量存储模型
_BERT_MODEL = AutoModel.from_pretrained('E:/bert-base-chinese').to('cuda' if torch.cuda.is_available() else 'cpu')
_BERT_TOKENIZER = AutoTokenizer.from_pretrained('E:/bert-base-chinese')
_TOPIC_KEYWORDS_MODEL = AutoModelForCausalLM.from_pretrained('E:/qianwen').to('cuda' if torch.cuda.is_available() else 'cpu')
_TOPIC_KEYWORDS_TOKENIZER = AutoTokenizer.from_pretrained('E:/qianwen')

# 定义加载文本文件的函数
# file_dir: 文本文件所在目录
# num_files: 要加载的文件数量，默认为 None 表示加载所有文件
def load_texts(file_dir, num_files=None):
    """加载文本文件"""
    # 获取目录下所有以 .txt 结尾的文件路径
    txt_files = [os.path.join(file_dir, f) for f in os.listdir(file_dir) if f.endswith('.txt')]
    # 如果指定了文件数量，则截取前 num_files 个文件
    if num_files is not None:
        txt_files = txt_files[:num_files]
    # 初始化存储文本的列表
    texts = []
    # 遍历每个文件路径
    for file_path in tqdm(txt_files, desc='加载文本文件'):
        # 以只读模式打开文件，指定编码为 utf-8
        with open(file_path, 'r', encoding='utf-8') as f:
            # 将文件内容添加到文本列表中
            texts.append(f.read())
    # 返回存储所有文本的列表
    return texts

# 定义获取文本嵌入的函数
# texts: 待处理的文本列表
def get_text_embeddings(texts):
    """使用 huggingface 预训练模型获取文本嵌入"""
    embeddings = []
    for text in tqdm(texts, desc='获取文本嵌入'):
        # 对文本进行分词并转换为 PyTorch 张量
        inputs = _BERT_TOKENIZER(text, padding=True, truncation=True, return_tensors='pt').to('cuda' if torch.cuda.is_available() else 'cpu')
        # 获取模型输出
        outputs = _BERT_MODEL(**inputs)
        # 使用 [CLS] 标记的嵌入作为文本嵌入
        embedding = outputs.last_hidden_state[:, 0, :]
        # 转换为 numpy 数组
        embedding = embedding.detach().cpu().numpy()
        embeddings.append(embedding)
    return np.vstack(embeddings)

# 定义保存文本嵌入结果的函数
# embeddings: 待保存的文本嵌入结果
# save_path: 保存路径，默认为 'embeddings.npy'
def save_embeddings(embeddings, save_path='embeddings.npy'):
    """保存文本嵌入结果"""
    # 使用 numpy 的 save 函数将嵌入结果保存到指定路径
    np.save(save_path, embeddings)

# 定义加载保存的文本嵌入结果的函数
# load_path: 加载路径，默认为 'embeddings.npy'
def load_embeddings(load_path='embeddings.npy'):
    """加载保存的文本嵌入结果"""
    # 使用 numpy 的 load 函数从指定路径加载嵌入结果并返回
    return np.load(load_path)

# 定义对文本嵌入进行聚类的函数
# embeddings: 待聚类的文本嵌入结果
# num_clusters: 聚类的数量，默认为 6
def cluster_texts(embeddings, num_clusters=6):
    """使用 KMeans 对文本嵌入进行聚类"""
    # 初始化 KMeans 聚类器，指定聚类数量和随机种子
    kmeans = KMeans(n_clusters=num_clusters, random_state=42)
    # 使用 KMeans 对嵌入结果进行拟合和预测，得到聚类标签
    labels = kmeans.fit_predict(embeddings)
    # 返回聚类标签
    return labels

# 定义获取每个聚类中前几篇文档的函数
# texts: 原始文本列表
# labels: 聚类标签列表
# num_clusters: 聚类的数量，默认为 6
# top_n: 每个聚类中要获取的文档数量，默认为 3
def get_top_documents_per_cluster(texts, labels, num_clusters=6, top_n=3):
    """获取每个聚类中的前几篇文档"""
    # 初始化存储每个聚类文档的字典
    cluster_docs = {i: [] for i in range(num_clusters)}
    # 遍历每个文档的索引和对应的聚类标签
    for idx, label in enumerate(labels):
        # 将文档添加到对应的聚类文档列表中
        cluster_docs[label].append(texts[idx])
    # 初始化存储每个聚类前几篇文档的列表
    top_docs = []
    # 遍历每个聚类
    for i in range(num_clusters):
        # 将每个聚类的前 top_n 篇文档添加到结果列表中
        top_docs.append(cluster_docs[i][:top_n])
    # 返回每个聚类的前几篇文档列表
    return top_docs

# 定义获取每个聚类主题词的函数
# top_docs_per_cluster: 每个聚类的前几篇文档列表

def get_topic_keywords(top_docs_per_cluster):
    """调用千问大语言模型获取每个聚类的主题词"""
    # 声明使用全局变量
    global _TOPIC_KEYWORDS_MODEL, _TOPIC_KEYWORDS_TOKENIZER

    # 初始化存储主题词的列表
    topic_keywords = []
    # 遍历每个聚类的前几篇文档
    for docs in tqdm(top_docs_per_cluster, desc='获取主题词'):
        # 构造对话消息
        # 截取每篇文档末尾1000字
        truncated_docs = [doc[-2000:] if len(doc) > 2000 else doc for doc in docs]
        messages = [
            {"role": "system", "content": "你是一个善于提炼文章关键词的作家，请从给定文档中提炼出主题词。"},
            {"role": "user", "content": "以下是同一类别的三篇文档，请给出它们的主题词：\n" + "\n\n".join(truncated_docs)}
        ]
        # 使用 apply_chat_template 构造输入文本
        text = _TOPIC_KEYWORDS_TOKENIZER.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        # 对输入文本进行分词
        model_inputs = _TOPIC_KEYWORDS_TOKENIZER([text], return_tensors='pt').to(_TOPIC_KEYWORDS_MODEL.device)
        # 使用模型获取结果
        generated_ids = _TOPIC_KEYWORDS_MODEL.generate(
            **model_inputs,
            max_new_tokens=50
        )
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]
        # 解码生成的文本
        result = _TOPIC_KEYWORDS_TOKENIZER.batch_decode(generated_ids, skip_special_tokens=True)[0]
        topic_keywords.append(result)
    # 返回每个聚类的主题词列表
    return topic_keywords

if __name__ == '__main__':
    # 文本文件目录，需根据实际情况修改
    file_dir = "E:/jjsw/jjsw-new"
    # 聚类的数量
    num_clusters = 20
    # 每个聚类中要获取的文档数量
    top_n = 3
    # 处理的文件数量，设为 None 则处理所有文件
    num_files = 71131
    # 文本嵌入结果的保存路径
    embeddings_save_path = 'embeddings.npy'

    # 加载文本
    texts = load_texts(file_dir, num_files)
    
    # 检查 embeddings.npy 是否存在，若存在则直接加载，否则获取并保存文本嵌入
    if os.path.exists(embeddings_save_path):
        embeddings = load_embeddings(embeddings_save_path)
    else:
        embeddings = get_text_embeddings(texts)
        save_embeddings(embeddings, embeddings_save_path)
    # 聚类
    labels = cluster_texts(embeddings, num_clusters)
    # 获取每个聚类的前几篇文档
    top_docs_per_cluster = get_top_documents_per_cluster(texts, labels, num_clusters, top_n)
    # 获取主题词
    topic_keywords = get_topic_keywords(top_docs_per_cluster)

    # 遍历每个聚类的主题词
    for i, keywords in enumerate(topic_keywords):
        # 打印每个聚类的主题词
        print(f'聚类 {i} 的主题词: {keywords}')
        
        # 创建以主题关键词命名的簇文件夹
        cluster_dir = os.path.join(os.getcwd(), keywords.strip().replace("\n",""))
        # 创建文件夹，如果已存在则忽略
        os.makedirs(cluster_dir, exist_ok=True)
        
        # 创建 xlsx 文件
        xlsx_path = os.path.join(cluster_dir, f'cluster_{i}_keywords.xlsx')
        # 初始化一个新的 Excel 工作簿
        wb = openpyxl.Workbook()
        # 获取工作簿的活动工作表
        ws = wb.active
        
        # 写入主题关键词
        ws.cell(row=1, column=1, value=keywords)
        
        # 获取当前簇的文档索引
        cluster_indices = [idx for idx, label in enumerate(labels) if label == i]
        # 获取目录下所有以 .txt 结尾的文件路径
        txt_files = [os.path.join(file_dir, f) for f in os.listdir(file_dir) if f.endswith('.txt')]
        # 如果指定了文件数量，则截取前 num_files 个文件
        if num_files is not None:
            txt_files = txt_files[:num_files]
        
        # 写入文档信息
        for row_num, idx in enumerate(cluster_indices, start=2):
            # 写入文档文件名
            ws.cell(row=row_num, column=1, value=os.path.basename(txt_files[idx]))
            # 写入文档内容
            ws.cell(row=row_num, column=2, value=texts[idx])
        
        # 保存 Excel 文件
        wb.save(xlsx_path)
